// Secure Socket wrapper using OpenSSL
// Copyright 2005 (C) Ralph Thomas

#include <util/string.h>
#include <util/socket.h>
#include <util/log.h>

#ifdef UNIX
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <unistd.h>
#include <poll.h>
#ifndef sockaddr_in
#include <netinet/in.h>
#endif
#endif

#ifdef _WINDOWS
#include <winsock2.h>
#define socklen_t int
#endif

#include <openssl/ssl.h>

namespace util {
	//
	/// This is an implementation of the socket class using plain BSD
	/// sockets. This implementation uses blocking sockets for
	/// simplicity's sake.
	//
	class socketSecureImpl : public socket {
	  private:
		int _fd;
		SSL_CTX* _context;
		SSL* _connection;
		bool _usingTLS;
	  public:
		socketSecureImpl() : _fd( -1 ), _context( 0 ), _connection( 0 ),
				  _usingTLS( false ) {
			static bool sslInit = false;
			//
			// Make sure that we initialize OpenSSL when starting.
			//
			if( !sslInit ) {
				SSL_load_error_strings();
				SSL_library_init();
				//
				// We should initialize the PRNG at this point
				// on most OSs it will initialize itself.
				//
				sslInit = true;
			}
			//
			// Create a context which we will use when we connect
			// to some server. We only allow ourselves to be a
			// TLS client -- we do not permit SSL connections.
			//
			_context = SSL_CTX_new( TLSv1_client_method() );
		}
		~socketSecureImpl() {
			if( _connection ) {
				SSL_shutdown( _connection );
				SSL_free( _connection );
			}
			if( _context ) SSL_CTX_free( _context );
			if( _fd != -1 ) {
#ifdef _WINDOWS
				closesocket( _fd );
				WSACleanup();
#else
				shutdown( _fd, SHUT_RDWR );
				close( _fd );
#endif
			}
		}

		bool connect( util::string hostname, int port ) {
			struct sockaddr_in servaddr;
			struct hostent* host;
#ifdef _WINDOWS
			//
			// WinSock initialization. This was inspired by
			// the Mozilla XPNet socket code.
			//
			int error;
			WORD winSockVersion;
			WSADATA winSockData;

			//
			// Try all of the different versions of winsock which
			// do what we need.
			//
			winSockVersion = MAKEWORD( 2, 2 );
			error = WSAStartup( winSockVersion, &winSockData );
			if( error == WSAVERNOTSUPPORTED ) {
				//
				// Try version 1.1
				//
				winSockVersion = MAKEWORD( 1, 1 );
				error = WSAStartup( winSockVersion, &winSockData );
				if( error == WSAVERNOTSUPPORTED ) return false;
			}
#endif
			_fd = ::socket( AF_INET, SOCK_STREAM, 0 );
			if( _fd == -1 ) return false;

			memset( &servaddr, 0, sizeof( servaddr ) );
			servaddr.sin_family = AF_INET;
			servaddr.sin_port = htons( port );
			
			//XXX: where does this memory come from? do i free it?
			host = gethostbyname( hostname.c_str() );
			if( !host ) return false;

			memcpy( &servaddr.sin_addr, (struct in_addr**)
				host->h_addr_list[0],
				sizeof( struct in_addr ) );

			if( ::connect( _fd, (struct sockaddr*) &servaddr,
				sizeof( servaddr ) ) < 0 ) return false;
			//
			// We can now create the SSL object from the context
			// we made. This SSL object is a "session" which is
			// associated with the context.
			//
			_connection = SSL_new( _context );
			if( !_connection ) return false;
			//
			// Attach the socket to OpenSSL.
			//
			if( !SSL_set_fd( _connection, _fd ) ) return false;
			//
			// All done for now, we use read and write on the
			// socket until startTLS is called, then we use
			// SSL_read and SSL_write.
			//
			return true;
		}

		virtual errorCode send( util::string message ) {
			char* buf;
			const int bufSize = message.size() + 1;
			int rc = -1;

			if( _fd == -1 ) return kDisconnected;

			buf = new char[ bufSize ];
			for( unsigned int i = 0; i < message.size(); i++ )
				buf[i] = message[i];
			buf[ message.size() ] = 0;
			//
			// If we're not using TLS yet, then we can use the
			// standard (BSD) send call.
			//
			if( !_usingTLS )
				rc = ::send( _fd, buf, message.size(), 0 );
			else rc = SSL_write( _connection, buf, message.size() );
			//
			// Clean up our allocated buffer and return the
			// appropriate code.
			//
			delete[] buf;
			if( rc <= 0 ) return kFailure;
			return kOk;
		}

		virtual errorCode recieve( util::string& message, int nbytes ) {
			const int bufSize = nbytes + 1;
			char* buf = new char[ bufSize ];
			util::string tmpMsg;
			int rc = -1;
			//
			// Make sure that we have an open connection to
			// recieve from.
			//
			if( _fd == -1 ) return kDisconnected;
			for( int i = 0; i < bufSize; i++ ) buf[i] = 0;
			//
			// If we're not in a TLS session yet, then we can use
			// the standard recv call.
			//
			if( !_usingTLS ) rc = recv( _fd, buf, nbytes, 0 );
			else rc = SSL_read( _connection, buf, nbytes );
			//
			// Clean up our buffer and return the appropriate
			// error code.
			//
			if( rc < 1 ) { delete[] buf; return kFailure; }
			message = buf;
			delete[] buf;
			return kOk;
		}

		virtual errorCode recieveLine( util::string& message ) {
			util::string line;
			util::string b;

			while( recieve( b, 1 ) == kOk ) {
				line += b;
				if( b == "\n" ) { message = line; return kOk; }
			}
			return kFailure;
		}

		virtual bool isDataAvailable( int timeout ) {
			fd_set socketSet;
			struct timeval socketTime;
			//
			// Regardless of if we're using TLS or not, we still
			// need an open file descriptor.
			//
			if( _fd == -1 ) return false;
			//
			// We have to use the SSL_pending call if we're using
			// SSL right now.
			//
			if( _usingTLS ) {
				//
				// If there is something already decoded in
				// the buffer then we return true, otherwise
				// we just select on the socket to see if any
				// encrypted data comes through. This is better
				// than the alternative, which would be to 
				// sleep for timeout if there was nothing here.
				//
				if( SSL_pending( _connection ) > 0 )
					return true;
			}
			//
			// Okay, we can just select on the socket, then.
			//
			FD_ZERO( &socketSet );
			FD_SET( _fd, &socketSet );
			socketTime.tv_sec = timeout;
			socketTime.tv_usec = 0;

			int rv = select( _fd+1, &socketSet, NULL, NULL,
					 &socketTime );
			if( (rv == 0) || (rv == -1) ) return false;
			return true;
		}

		virtual bool startTLS() {
			if( !_connection ) return false;
			if( SSL_connect( _connection ) == 1 ) {
				_usingTLS = true;
				return true;
			}
			return false;
		}
	};

	socket* socket::fromHostnameAndPort( util::string hostname, int port ) {
		socketSecureImpl* _socket = new socketSecureImpl();
		if( _socket->connect( hostname, port ) ) return _socket;
		delete _socket;
		return NULL;
	}
	
	socket::~socket() {}
};
